'''------------------------------------------------------------------------------------------------
 $ @Author: Ysq
 $ @Date: 2021-12-23 19:16:28
 $ @LastEditTime: 2021-12-23 22:11:40
 $ @LastEditors: Please set LastEditors
 $ @Description: 一个轻量化的，跟踪结果输出器
 $ @FilePath: \YSQTrackerEvaluationToolKit_v1.0\EvalToolKit\outPutVisRes.py
--------------------------------------------------------------------------------------------------'''
import sys

sys.path.append('./')


from GetCfg.getCfg_OutPutVisRes import getYamlCfg
from GetCfg.setParam_OutPutVisRes import getParam
from tqdm import tqdm
from EvalToolKit.evalToolKit import EvalToolKit
from EvalToolKit.tracker import Tracker
from Tools.MTree import MTree

class OutPutVisRes(EvalToolKit):
    def __init__(self) -> None:
        super().__init__()
        self.args = getParam()
        self.YamlCfg = getYamlCfg(self.args.videosNameListDir)

    def init(self):
        print(MTree(0,'初始化 ToolKit'))
        # -read anno
        # ---- read ground truth
        for name in tqdm(self.YamlCfg.outputVideosList,desc=MTree(1,'Load Ground Truth')):
            if name in self.args.starFrameAndEndFrame.keys():
                if self.args.starFrameAndEndFrame[name][1] == 0:
                    self.gt[name] = self.loadAnno(self.args.testSetDir+'/'+name+'/groundtruth_rect.txt')[
                        self.args.starFrameAndEndFrame[name][0]:]
                else:
                    self.gt[name] = self.loadAnno(self.args.testSetDir+'/'+name+'/groundtruth_rect.txt')[
                        self.args.starFrameAndEndFrame[name][0]:self.args.starFrameAndEndFrame[name][1]]
            self.gt[name] = self.loadAnno(self.args.testSetDir+'/'+name+'/groundtruth_rect.txt')
        # ---- 读每个跟踪器的跟踪结果，并初始化每个跟踪器的 tracker 类
        colorIdx = 0
        for tracker in self.args.trackerBboxDir.keys():
            tempAnnos = {}
            for i in tqdm(range(len(self.YamlCfg.outputVideosList)),desc = MTree(1,'Load '+tracker+' Bbox')):
                name = self.YamlCfg.outputVideosList[i]
                if name in self.args.starFrameAndEndFrame.keys():
                    if self.args.starFrameAndEndFrame[name][1] == 0:
                        tempAnnos[name] = self.loadAnno(self.args.trackerBboxDir[tracker]+'/'+name+'.txt')[
                            self.args.starFrameAndEndFrame[name][0]:]
                    else:
                        tempAnnos[name] = self.loadAnno(self.args.trackerBboxDir[tracker]+'/'+name+'.txt')[
                            self.args.starFrameAndEndFrame[name][0]:self.args.starFrameAndEndFrame[name][1]]
                tempAnnos[name] = self.loadAnno(self.args.trackerBboxDir[tracker]+'/'+name+'.txt')
            self.tracker[tracker] = Tracker(tracker,self.args.trackerBboxDir[tracker],self.color[colorIdx],tempAnnos)
            colorIdx += 1
        self._checkOutputFolder()